#!/usr/bin/env python3
# -*-coding:utf8-*-

import json
import logging
import os
import time
from pathlib import Path
from typing import Dict, List, Sequence

import numpy as np
from piper_sdk import C_PiperInterface_V2
from scipy.spatial.transform import Rotation

from sim.mj_ik import MujocoIK

# from model.transforms import normalize_end, denormalize_joint, normalize_joint


log = logging.getLogger(__name__)


def arm_end_translation(xyz: Sequence[float], rpy: Sequence[float], shift_len: float):
    """
    Given the arm position and rpy, estimate the end-effector position
    xyz: mm
    rpy: radian
    shift_len: mm

    return: mm
    """
    # Use "x,y,z": external rotation axis. NOT "XYZ"
    R = Rotation.from_euler("xyz", rpy, degrees=False)
    shift = R.apply(np.array([0, 0, shift_len]))
    return xyz + shift


def to_machine_angle(joints: np.ndarray):
    """
    radian -> 0.001 degree
    """
    return (np.rad2deg(joints) * 1000).astype(int)


def from_machine_angle(joints: np.ndarray):
    """
    0.001 degree -> radian
    """
    return np.deg2rad(joints) / 1000


def clip_joint_radian(joints: np.ndarray):
    """
    clip the joint radian to the working range of the robot
    """
    max_rad = np.array([2.6179, 3.14, 0, 1.745, 1.22, 2.09439])
    min_rad = np.array([-2.6179, 0, -2.967, -1.745, -1.22, -2.09439])
    return np.clip(joints, min_rad, max_rad)


class PiperArm:
    """
    Unify the measurement unit:
    Length: mm
    Angle: radian
    Time: s
    Speed: mm/s

    __speed: mm/s (int) gives the running speed of the arm
    __Hz: Hz (int) gives the frequency of the signal sent to the arm
    __timeout: s (int) gives the timeout of the arm execution
    __gripper_effort: 0~5000, 0~5N/m
    __gripper_angle: 0~10000, 0~100°

    gripper_length (mm): is used for estimating the arm position from the end effector position
    """

    __speed: int = 40
    __Hz: int = 100
    __timeout: int = 5
    __gripper_effort: int = 1000
    __wait_gripper: float = 0.3

    def __init__(
        self,
        can: str = "can0",
        ik_model=None,
        planning_model=None,
        wait_execution: bool = True,
        gripper_length: float = 140,
        execution_horizon: int = 1,
    ):  # in mm

        self.__timestep = 1 / self.__Hz
        self.can = can
        self.wait_execution = wait_execution
        self.gripper_length = gripper_length

        # Initialize the robot
        self.arm = C_PiperInterface_V2(self.can)
        self.arm.ConnectPort()
        if self.reset():
            print(f"Sucessfully connect to the robot CAN={can}")
            print(f"Sucessfully make arm {can} enabled")

        # Initialize the models
        self.ik_model = ik_model
        self.planning_model = planning_model

        # initialize the targets position and orientation of the effector
        self.ik = MujocoIK()
        self.execution_horizon = execution_horizon
        # status = self.get_effector_status()

        # self.target_pos = status["xyz"] / 1000  # mm -> m
        # self.target_quat = Rotation.from_euler("xyz", status["rpy"]).as_quat()[[3, 0, 1, 2]]  # (w, x, y, z)
        self.tolerance = 0.001  # 1mm tolerance
        self.quat_prior = np.array([0.4, 0, 0.8, 0])  # (w, x, y, z)

    def reset(self):
        s_time = time.time()
        self.go_zero()
        while not all(self.arm.GetArmEnableStatus()):
            self.arm.EnableArm(7)
            self.go_zero()
            time.sleep(self.__timestep)
            if time.time() - s_time > self.__timeout:
                log.error("Arm reset timeout!")
                return False
        return True

    def go_zero(self):
        """
        use joint mode to move to the zero position
        """
        self.arm.ModeCtrl(0x01, 0x01, self.__speed, 0x00)
        self.arm.JointCtrl(0, 0, 0, 0, 0, 0)
        self.arm.GripperCtrl(0, self.__gripper_effort, 0x01, 0)
        if self.wait_execution:
            self.wait_motion_done()

    def set_speed(self, speed: int):
        self.__speed = speed

    def planning_to(self, pos: np.ndarray, quat: np.ndarray = None):
        if quat is None:
            quat = self.quat_prior

        assert self.planning_model is not None, "planning_model is not set"
        current_status = self.get_joint_status()
        j0 = current_status["joint"]

        pos = pos / 1000  # mm -> m
        self.ik.move_target_to(pos, quat)
        if self.ik.effector_target_dist() < self.tolerance:
            print(f"Target reached! dist={self.ik.effector_target_dist()}m < {self.tolerance}m")
            return True

        cond = np.concat([j0, pos, quat], axis=-1)
        # (batch_size, num_act_samples, horizon_steps, action_dim)
        joints = self.planning_model(cond)[0]

        dists = [self.ik.effector_target_dist(j[self.execution_horizon - 1], pos, quat) for j in joints]

        for j in joints[np.argmin(dists), : self.execution_horizon]:
            self.move_to_joint(j)
            self.ik.step(np.concat([j, current_status["gripper"]]))
            if self.ik.effector_target_dist() < self.tolerance:
                print(f"Target reached! dist={self.ik.effector_target_dist()}m < {self.tolerance}m")
                return True

    def move_to(self, xyz: Sequence[float]):
        """
        USE THIS TO REACH CONTACT POINT !!!!

        move to the contact position (gripper)
        xyz (mm) : the position of the end-effector

        The process:
        1. clip the xyz to the same zone
        2. If ik model loaded: predict the joint state -> joint move
        3. If no ik model: estimate the arm position -> arm move
        """
        xyz = self.safety_clip(xyz)
        if self.ik_model is not None:
            # remove batch dimension
            pred_joint = self.ik_model(xyz)[0]
            return self.move_to_joint(pred_joint)
        arm_xyz, arm_rpy = self.est_arm_pos_from_ee(xyz)
        return self._move_arm_to(arm_xyz, arm_rpy)

    def ik_render(self):
        status = self.get_joint_status()
        j = np.concat([status["joint"], status["gripper"]])
        self.ik.step(j)
        self.ik.render()

    def open_gripper(self, width: int = 100):
        """
        width: 0~100
        need some time to finish the operation: self.__wait_gripper
        """
        self.wait_motion_done()
        self.arm.GripperCtrl(width * 1000, self.__gripper_effort, 0x01, 0)
        time.sleep(self.__wait_gripper)

    def close_gripper(self):
        """
        need some time to finish the operation: self.__wait_gripper
        """
        self.wait_motion_done()
        self.arm.GripperCtrl(0, self.__gripper_effort, 0x01, 0)
        time.sleep(self.__wait_gripper)

    def move_to_joint(self, joints: np.ndarray):
        """
        joints: radian -> clip -> 0.001 degree (int) -> machine execution
        """
        assert len(joints) == 6, "Numer of target joint status is 6"
        joints = clip_joint_radian(joints)
        machine_joints = to_machine_angle(joints)

        # "Joint" mode
        self.arm.ModeCtrl(0x01, 0x01, self.__speed, 0x00)
        self.arm.JointCtrl(*machine_joints)
        if self.wait_execution:
            return self.wait_motion_done()

    def _move_arm_to(self, xyz: np.ndarray, rpy: np.ndarray):
        """
        Move the arm to the given position and orientation of the ARM.
        It is NOT the end-effector position.

        xyz: mm
        rpy: radian
        """
        # convert to machine units (int)
        xyz = (xyz * 1000).astype(int).tolist()
        rpy = to_machine_angle(rpy).tolist()

        # "Position" mode
        self.arm.ModeCtrl(0x01, 0x00, self.__speed, 0x00)
        self.arm.EndPoseCtrl(*(xyz + rpy))
        if self.wait_execution:
            return self.wait_motion_done()

    def get_arm_pos_status(
        self,
    ) -> Dict:
        arm_pos = self.arm.GetArmEndPoseMsgs().end_pose
        xyz = np.array([arm_pos.X_axis, arm_pos.Y_axis, arm_pos.Z_axis]) / 1000
        rpy = from_machine_angle(np.array([arm_pos.RX_axis, arm_pos.RY_axis, arm_pos.RZ_axis]))
        gripper = from_machine_angle(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle).item() / 1.78 * 0.024
        return {
            "xyz": xyz,  # in mm
            "rpy": rpy,  # in radian
            "gripper": np.array([1, -1]) * gripper,
        }

    def get_joint_status(
        self,
    ) -> Dict:
        arm_status = self.arm.GetArmJointMsgs().joint_state
        jo = [
            arm_status.joint_1,
            arm_status.joint_2,
            arm_status.joint_3,
            arm_status.joint_4,
            arm_status.joint_5,
            arm_status.joint_6,
        ]
        gripper = from_machine_angle(self.arm.GetArmGripperMsgs().gripper_state.grippers_angle).item() / 1.78 * 0.024
        return {
            "joint": from_machine_angle(np.asarray(jo)),
            "gripper": np.array([1, -1]) * gripper,
        }

    def get_effector_status(self):
        status = self.get_arm_pos_status()
        ee_xyz = arm_end_translation(status["xyz"], status["rpy"], self.gripper_length)
        return {
            "xyz": ee_xyz,
            "rpy": status["rpy"],
            "gripper": status["gripper"],
        }

    def est_arm_pos_from_ee(self, xyz: Sequence[float], ry: float = 0.01):
        """
        A manual function to estimate the arm xyz&rpy from the end-effector position.
        It only works for the arm near the table. The estimated position signal is
        processed by the internal inverse kinematics model,
        which often fails.
        xyz: mm
        ry: radian
        """
        # ry \approx the angle from gripper's vertical postion
        x, y, z = xyz
        joint1 = np.arctan(y / x)
        rx = np.sign(joint1) * np.pi
        rz = joint1 - np.pi * np.sign(joint1)

        xy_angle = np.arctan(y / x)
        x = x - self.gripper_length * np.sin(ry) * np.cos(xy_angle)
        y = y - self.gripper_length * np.sin(ry) * np.sin(xy_angle)
        z = z + self.gripper_length * np.cos(ry)

        return np.array([x, y, z]), np.array([rx, ry, rz])

    def is_in_motion(self):
        return bool(self.arm.GetArmStatus().arm_status.motion_status)

    def wait_motion_done(self):
        start_t = time.time()
        time.sleep(self.__timestep)
        while self.is_in_motion():
            time.sleep(self.__timestep)
            if time.time() - start_t > self.__timeout:
                log.error(f"Excution exceeds timeout threshold={self.__timeout}s")
                return False
        return True

    def is_gripping_sth(self):
        return self.arm.GetArmGripperMsgs().gripper_state.grippers_effort < -200

    def follow_traj(self, filename: Path):
        """
        Only support joint mode
        """
        with open(filename, "r") as f:
            status = json.load(f)
        for p_ in status["joint"]:
            self.move_to_joint(p_)
            time.sleep(self.__timestep)

    def record_traj(self, filename: Path = "data/json_data/trajectories.json"):
        xyzs, rpys, joints, lengths = [], [], [], []
        try:
            while True:
                xyz, rpy, j = [], [], []
                while self.arm.GetArmStatus().arm_status.teach_status == 1:
                    status = self.get_effector_status()
                    xyz.append(status["xyz"].tolist())
                    rpy.append(status["rpy"].tolist())
                    j.append(self.get_joint_status()["joint"].tolist())
                    time.sleep(0.1)
                    print(f"Recorded {len(xyz)} points", end="\r")

                if xyz and rpy and j:
                    xyzs += xyz
                    rpys += rpy
                    joints += j
                    lengths.append(len(xyz))
                    print(f"Recorded {len(xyz)} points: Done!", end="\n")
                else:
                    print("No points recorded, please move the arm to record...", end="\r")
                    time.sleep(0.1)
        except KeyboardInterrupt:
            with open(filename, "w") as f:
                json.dump(
                    {
                        "xyz": xyzs,
                        "rpy": rpys,
                        "joint": joints,
                        "length": lengths,
                    },
                    f,
                )

    def record_points(
        self, save_dir: Path = "data/json_data", method: str = "auto", tag: str = "", time_interval: float = 0.02
    ):
        """
        record the joint and ee position status
        """
        n_ = 0
        joints, xyz, rpy = [], [], []

        if not os.path.exists(save_dir):
            os.mkdir(save_dir)

        try:
            while True:
                if method != "auto":
                    key = input("Place the gripper and press ENTER...")
                    if key == "q":
                        assert 1 == 0
                joints.append(self.get_joint_status()["joint"].tolist())
                status = self.get_effector_status()
                xyz.append(status["xyz"].tolist())
                rpy.append(status["rpy"].tolist())
                n_ += 1

                time.sleep(time_interval)
                print(f"Saved {n_} points", end="\r")

        except KeyboardInterrupt:

            f_name = os.path.join(save_dir, f"points_{tag}.json")
            with open(f_name, "w") as f:
                json.dump(
                    {
                        "joint": joints,
                        "xyz": xyz,
                        "rpy": rpy,
                    },
                    f,
                )
                print(f"{n_} data points are saved to {f_name}")

    @staticmethod
    def safety_clip(xyz: Sequence[float]):
        return np.array(xyz)

    def mujoco_viewer(self, run_in_background=True):
        """
        Used for notebook running
        """
        os.environ["MUJOCO_GL"] = "egl"
        import threading

        import mujoco
        import mujoco.viewer

        model = mujoco.MjModel.from_xml_path("piper_description/mujoco_model/piper_description.xml")
        data = mujoco.MjData(model)

        def _viewer_loop():
            print("Viewer launced!")
            with mujoco.viewer.launch_passive(model, data) as viewer:
                print("Viewer launched!")
                while viewer.is_running():
                    joint = np.concat(
                        [self.get_joint_status()["joint"], np.array([1, -1]) * self.get_joint_status()["gripper"]]
                    )
                    data.qpos = joint
                    mujoco.mj_step(model, data)
                    viewer.sync()
                    time.sleep(0.01)
                print("Viewer closed!")

        if run_in_background:
            # daemon=True: thread will be terminated when the main thread terminates
            self.viewer_thread = threading.Thread(target=_viewer_loop, daemon=True)
            self.viewer_thread.start()
            print("Viewer started in background thread. You can now run other cells!")
            return self.viewer_thread
        else:
            _viewer_loop()


if __name__ == "__main__":

    def test_planning():
        import hydra
        from omegaconf import DictConfig

        with hydra.initialize(version_base=None, config_path="cfg"):
            ej_cfg: DictConfig = hydra.compose(
                config_name="train_diffusion_planning",
                overrides=[
                    "tag=partial",
                    "model.temperature=0.1",
                    "model.horizon_steps=10",
                    "model.num_act_samples=100",
                ],
            )
            model = hydra.utils.instantiate(ej_cfg.model)
            model.load_ckpt(ej_cfg.ckp_name + ".ckp")

        self = PiperArm("can0", wait_execution=False, planning_model=model, execution_horizon=10)
        pos = np.array([408.30336375, 68.6895842, 45.59737271])
        quat = np.array([-0.28161, 0.08522, -0.95511, -0.03461])
        # for i in range(10):
        self.planning_to(pos=pos, quat=quat)

    # from model import *
    # from model.diffusion.diffusion_ctrl import *

    # # ej_model = DiffusionCtrlModel(
    # #     in_dim=3,
    # #     out_dim=6,
    # #     model_file="diffusion_end_joint_v0.ckp",  # "diffusion_end_joint_all_normed.ckp",
    # #     dropout_rate=0.1,
    # #     temperature=0.1,
    # #     fk_model=MLPCtrlModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp").model,
    # # )
    # # ej_model = DiffusionCtrlModel(
    # #     in_dim=3,
    # #     out_dim=6,
    # #     model_file="diffusion_end_joint_T10.ckp",  #
    # #     # dropout_rate=0.1,
    # #     denoising_steps=10,
    # #     # model_file="diffusion_end_joint_nodropout_v1.ckp",
    # #     temperature=0.1,
    # #     num_act_samples=100,
    # #     # fk_model=EnsembleFKModel(in_dim=6, out_dim=3, num_heads=10, model_file="ensemble_FK.ckp"),
    # #     fk_model=MLPCtrlModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp").model,
    # # )
    # # plan_model = DiffusionPlanningModel(
    # #     in_dim=12,
    # #     out_dim=6,
    # #     model_file="diffusion_planning_T10_H10_S1.ckp",  # "diffusion_end_joint_all_normed.ckp",
    # #     horizon_steps=10,
    # #     temperature=0.1,
    # #     denoising_steps=10,
    # #     num_act_samples=1,
    # #     # fk_model=MLPCtrlModel(in_dim=6, out_dim=3, model_file="all_normed_FK.ckp").model,
    # # )
    # # self = PiperArm("can0", wait_execution=False, end_joint_model=ej_model)
    # # self.planning_model = plan_model
    # # diffusion partial ik
    # ik_model = DiffusionIKModel(
    #     in_dim=3,
    #     out_dim=6,
    #     model_file="diffusion_partial_ik.ckp",
    #     temperature=0.1,
    #     num_act_samples=200,
    #     fk_model=MlpPartialFKModel(in_dim=6, out_dim=3, model_file="mlp_partial_fk.ckp").model,
    # )

    # # ik_model = MlpPartialIKModel(model_file="mlp_partial_ik.ckp")
    # self = PiperArm("can0", wait_execution=False, ik_model=ik_model)

    # self.go_zero()
    # # p = np.array([435.32488388 , -24.77839234  , 1.09628048])
    # # p = np.array([401.21426989, 40.37996443,  5.94443308])
    # p = np.array([408.30336375, 68.6895842, 45.59737271])

    # # p = np.array([ 55.59858803, -46.4059748 ,  -1.60010257]) ## danger!

    # # pp = np.array([10.532, 85.783, -25.349, 2.987, 21.791, -1.417])

    # # self.move_to_arm_pos(np.array([ 411.856,  -17.604,  130.043, -177.575,   11.987, -178.31 ]))
    # # [ 419.21315514 -109.77677174    3.67021936]
    # # self.move_to(p)
    # # self.planning_move_to(p)
    # self.go_zero()
